/**
  ETFAna project, Anyang Normal University && IMP-CAS
  \class ETFVTrack
  \brief class for a track. Note that it is supposed to be a virtual class
  \author SUN Yazhou, asia.rabbit@163.com
  \since 2021-11-01
  \date 2021-12-14 last modified
  \attention
  changelog
  <table>
  <tr>  <th>Date         <th>Author      <th>Description                   </tr>
  <tr>  <td>2021-11-01   <td>Asia Sun    <td>file created                  </tr>
  </table>

  \copyright Copyright (c) 2021-2024 Anyang Normal U. && IMP-CAS with LGPLv3 LICENSE
*/

#include <iostream>
#include <TGraph.h>
#include "ETFVTrack.h"
#include "ETFMsg.h"
#include "ETFMath.h"
#include "ETFDCHit.h"
#include "ETFDetector.h"
#include "ETFHPTDCCh.h"

using std::cout;
using std::endl;

ETFVTrack::ETFVTrack() : TObject(), fChiPD(-1.), fGraph(nullptr), fColor(4){
  Initialize();
} // end ctor

// copy ctor
ETFVTrack::ETFVTrack(const ETFVTrack &t) : TObject(t){
  *this = t;
} // end copy ctor

// assignment ctor
ETFVTrack &ETFVTrack::operator=(const ETFVTrack &t){
  if(&t == this) return *this;

  copy(&t);

  return *this;
} // end member function operator=

// the same as copy ctor
void ETFVTrack::copy(const ETFVTrack *t){
  fName = t->fName;
  fK = t->fK; fB = t->fB;
  fD2 = t->fD2; fChi2 = t->fChi2;
  fDT0 = t->fDT0; fNofHits = t->fNofHits;
  fTOTAvrg = t->fTOTAvrg;
  fFitStatus = t->fFitStatus;
  fChiPD = t->fChiPD;
  fHits = t->fHits;
  fHitPosMap = t->fHitPosMap;
  fLayerId = t->fLayerId;
  fColor = t->fColor;
} // end member function Copy

ETFVTrack::~ETFVTrack(){
  if(fGraph) delete fGraph;
} // end dtor

void ETFVTrack::Initialize(){
  fHits.clear(); fHitPosMap.clear();
  fLayerId.clear();
  fK = -999.; fB = -999.;
  fD2 = -999.; fChi2 = -999.;
  fDT0 = -999.; fNofHits = 0;
  fTOTAvrg = -999.;
  fFitStatus = FitStatus::kUNFITTED;
} // end member function Initialize

// set hit fired status for all hits
void ETFVTrack::SetHitFiredStatus(int s){
  for(auto &p : fHits) p.second->SetFiredStatus(s);
} // end member function SetHitFiredStatus

void ETFVTrack::AddHit(ETFDCHit *h){
  if(!h) ETFMsg::Error("ETFVTrack", "AddHit: input pointer is nullptr");
  fHits[h->GetLayerId()] = h;
  fLayerId.push_back(h->GetLayerId());
  fNofHits++;
} // end member function AddHit

void ETFVTrack::SetId(int id, short type){
  const char *xuv = "xuv";
  fName = ETFMsg::Form("vTrack%c_#%d", xuv[type%3], id);
} // end member function SetId

double ETFVTrack::DistanceToLine(double z, double x) const{
  if(-999. == fK || -999. == fB)
    ETFMsg::Error("ETFVTrack", "DistanceToLine: track not assigned");
  return ETFMath::DistanceToLine(z, x, fK, fB);
} // end member function DistanceToLine

// using external time as drift time start
// reset fFitStatus if drift time is updated
void ETFVTrack::SetDriftTimeStart(double time0){
  if(-999. == time0)
    ETFMsg::Error(GetName(), "SetDriftTimeStart: invalid drift time start: %f", time0);
  bool isNew = false;
  for(auto &p : fHits) if(p.second->SetDriftTimeStart(time0)) isNew = true;
  if(isNew || -999. == fDT0){
    fDT0 = time0;
    fFitStatus = FitStatus::kUNFITTED;
  } // end if
} // end member function SetDriftTimeStart

// update drift times for all member hits
void ETFVTrack::UpdateDriftTime(){
  SetDriftTimeStart(fDT0);
} // end member function UpdateDriftTime

// k in angle: atan(k)/DEGREE
double ETFVTrack::GetAng() const{
  return ETFMath::ang(fK);
} // end member function GetAng

// calculate dsquare: i.e. linear fit without drift distance info
double ETFVTrack::GetD2(){
  if(-999. != fD2) return fD2;

  if(!fHits.size()) ETFMsg::Error("ETFVTrack", "GetD2: track not assigned");
  const int n = fHits.size();
  vector<double> z, x, dx;
  z.resize(n); x.resize(n); dx.resize(n);
  dx.assign(n, -1.); // dx not unkown, so set to minus deliberately
  int i = 0;
  for(auto p : fHits){ p.second->GetProjZX(z[i], x[i]); i++; }
  ETFMath::LinearFit(z.data(), x.data(), dx.data(), n, fB, fK, fD2);

  return fD2;
} // end member function GetDsquare

// fit the track with drift distance info
// \retval true: if the fitting result meets user criteria; false: otherwise
bool ETFVTrack::Fit(){
  if(FitStatus::kUNFITTED != fFitStatus){
    if(FitStatus::kBADFIT == fFitStatus) return false;
    if(FitStatus::kGOODFIT == fFitStatus) return true;
  } // end if

  const int n = fHits.size();
  if(!n) ETFMsg::Error("ETFVTrack", "Fit: track not assigned");
  vector<double> z, x, r;
  z.resize(n); x.resize(n); r.resize(n);
  int i = 0;
  for(auto &p : fHits){ p.second->GetProjZXR(z[i], x[i], r[i]); i++; }

  // fit the track and assign slope and intercept
  // fChi2 = fD2;
  // for(i = 2; i--;) // iterated to get finer results
  fChi2 = ETFMath::IterFit(n, z.data(), x.data(), r.data(), fK, fB, fD2);
  UpdateChi(); // update chi so as to be used in GoodFit()

  if(GoodFit()){ fFitStatus = FitStatus::kGOODFIT; return true; }
  fFitStatus = FitStatus::kBADFIT;
  return false;
} // end member function Fit

void ETFVTrack::UpdateChi(){
  if(-999. == fK || -999. == fB)
    ETFMsg::Error("ETFVTrack", "UpdateChi: (k,b) not set");
  const double d = 1./sqrt(1.+fK*fK);
  fTOTAvrg = 0.;
  for(auto &p : fHits){
    auto &h = p.second;
    h->SetChi(fabs(fK*h->GetProjZ() - h->GetProjX() + fB)*d - h->GetR());
    fTOTAvrg += h->GetTChannel()->GetTOT();
  } // end for over hits
  fTOTAvrg /= fHits.size();
} // end member function UpdateChi

// test if a fit is a good fit (d2, chi and chi2)
bool ETFVTrack::GoodFit() const{
  const int n = fHits.size();
  if(fChi2 > fChiPD*fChiPD*n) return false;
  for(auto &p : fHits)
    if(fabs(p.second->GetChi()) > fChiPD)
      return false;
  return true;
} // end member function GoodFit

// \return whether layer with layerId is fired
bool ETFVTrack::fired(int layerId) const{
  for(const auto &i : fLayerId) if(i == layerId) return true;
  return false;
} // end member function fired


ETFDCHit *ETFVTrack::operator[](int layerId) const{
  if(fHits.count(layerId)) return fHits.at(layerId);
  return nullptr;
} // end member function operator[]

double ETFVTrack::GetHitPos(const string &d) const{
  if(!fHitPosMap.count(d))
    ETFMsg::Warn(GetName(), "GetHitPos: no record for det %s", d.c_str());
  return fHitPosMap.at(d);
} // end member function GetHitPos
void ETFVTrack::SetHitPos(const string &d, double x){
  fHitPosMap[d] = x;
} // end member function SetHitPos

// assign hit pos in lab frame (trk around TA)
void ETFVTrack::UpdateHitPos(ETFDetector *d){
  const string name = d->GetName();
  const double p1[2] = {GetK(), GetB()};
  double p[3]{};
  d->GetHitPosition(p1, p);
  fHitPosMap[name] = p[0];
} // end member function UpdateHitPos

void ETFVTrack::Draw(const char *){
  fGraph->Draw("plsame");
} // end member function Draw

// update fGraph in [z0, z1]
void ETFVTrack::ConstructGraph(double z0, double z1){
  if(!fGraph){
    fGraph = new TGraph();
    fGraph->SetLineColor(fColor);
    fGraph->SetLineWidth(2);
  } // end if
  while(fGraph->GetN()) fGraph->RemovePoint(0);
  fGraph->SetPoint(0, z0, fK*z0+fB);
  fGraph->SetPoint(1, z1, fK*z1+fB);
} // end member function ConstructGraph

void ETFVTrack::Print(const char *) const{
  cout << "\033[34;1m________Track ";
  cout << fName << " Print______________\033[0m" << endl;
  cout << "k: " << fK << " b: " << fB << endl;
  cout << "Number of Hits: " << GetNofHits() << endl;
  cout << "fired layer id:";
  for(const auto &l : fLayerId) cout << " " << l;
  cout << endl << "fDT0 (drift time start): " << fDT0 << " ns" << endl;
  cout << "d2: " << fD2 << " fChi2: " << fChi2 << endl;
  cout << "chi per layer: " << fChi2/GetNofHits() << endl;
  cout << "\033[33;1m  ~~~~~~ Hits at Detectors ~~~~~~~\033[0m" << endl;
  for(auto &p : fHitPosMap)
    cout << "Det: " << p.first << " lx: " << p.second << endl;
  cout << "\033[33;1m  ~~~~~~ All the Hits ~~~~~~~\033[0m" << endl;
  for(const auto &p : fHits) p.second->Print();
  cout << "\033[1m________End Track Print______________\033[0m" << endl;
} // end member function Print
